# Lorenz模型构造
import matplotlib.pyplot as plt
import numpy as np
from pylab import mpl
"""
# 洛伦兹方程
dx/dt = - a * X + a * Y
dy/dt = b * X - Y -X * Z 
dz/dt = c * Z + X * Y
"""

# 模型参数
a = 10
b = 28
c = 8/3
t = [0]
x = [1]
y = [0]
z = [0]
h = 0.01
n = 10000

mpl.rcParams['font.sans-serif'] = ['FangSong'] # 指定默认字体
mpl.rcParams['axes.unicode_minus'] = False # 解决保存图像是负号'-'显示为方块的问题

def Lorenz(x0, y0, z0, p, q, r, T):
    # 微分迭代步长
    h=0.01
    x=[]
    y=[]
    z=[]
    for t in range(T):
        xt = x0 + h * p * (y0 - x0)
        yt = y0 + h * (q * x0 - y0 - x0 * z0)
        zt = z0 + h * (x0 * y0 - r * z0)
        # x0、y0、z0统一更新
        x0, y0, z0 = xt, yt, zt
        x.append(x0)
        y.append(y0)
        z.append(z0)
    return x, y, z


# 辅助龙格库塔方程组进行迭代构造全部数据
# dx/dt = - a * X + a * Y
def FX(X, Y, Z):
    return - a * (X - Y)


# dy/dt = b * X - Y -X * Z
def FY(X, Y, Z):
    return b * X - Y - X * Z


# dz/dt = c * Z + X * Y
def FZ(X, Y, Z):
    return - c * Z + X * Y


# 利用龙格库塔方程组对Lorenz微方程组求解
def main():
    for i in range(n):
        K1 = FX(x[-1], y[-1], z[-1])
        L1 = FY(x[-1], y[-1], z[-1])
        M1 = FZ(x[-1], y[-1], z[-1])
        K2 = FX(x[-1] + h * K1 / 2, y[-1] + h * L1 / 2, z[-1] + h * M1 / 2)
        L2 = FY(x[-1] + h * K1 / 2, y[-1] + h * L1 / 2, z[-1] + h * M1 / 2)
        M2 = FZ(x[-1] + h * K1 / 2, y[-1] + h * L1 / 2, z[-1] + h * M1 / 2)
        K3 = FX(x[-1] + h * K2 / 2, y[-1] + h * L2 / 2, z[-1] + h * M2 / 2)
        L3 = FY(x[-1] + h * K2 / 2, y[-1] + h * L2 / 2, z[-1] + h * M2 / 2)
        M3 = FZ(x[-1] + h * K2 / 2, y[-1] + h * L2 / 2, z[-1] + h * M2 / 2)
        K4 = FX(x[-1] + h * K3, y[-1] + h * L3, z[-1] + h * M3)
        L4 = FY(x[-1] + h * K3, y[-1] + h * L3, z[-1] + h * M3)
        M4 = FZ(x[-1] + h * K3, y[-1] + h * L3, z[-1] + h * M3)
        x.append(x[-1] + h / 6 * (K1 + 2 * K2 + 2 * K3 + K4))
        y.append(y[-1] + h / 6 * (L1 + 2 * L2 + 2 * L3 + L4))
        z.append(z[-1] + h / 6 * (M1 + 2 * M2 + 2 * M3 + M4))


    # 构造Lorenz模型三维演化轨迹
    fig = plt.figure()
    ax = fig.add_subplot(projection='3d')
    ax.plot(x, y, z, label='parametric curve', c='green')

    plt.title('Lorenz模型三维演化轨迹')
    plt.xlabel('x轴')
    plt.ylabel('y轴')
    #plt.savefig('Lorenz三维演化轨迹.png')
    plt.show()


    # 给出迭代前的初值
    x0 = 1
    y0 = 0
    z0 = 0
    x1, y1, z1 = Lorenz(x0, y0, z0, a, b, c, n)
    # 初值进行微小的变化
    x0 = 1
    y0 = 0
    z0 = 0.00001
    xx, yy, zz = Lorenz(x0, y0, z0, a, b, c, n)
    t = np.arange(0, n)

    # x(t)轨道演化图
    plt.scatter(t, x1, s=1, c="green")
    plt.scatter(t, xx, s=1,)
    plt.title('x(t)轨道演化图')
    plt.xlabel('x轴')
    plt.ylabel('y轴')
    plt.savefig('Lorenz的x(t)轨道演化图.png')

    # y(t)轨道演化图
    #plt.scatter(t, y1, s=1, c='pink')
    #plt.scatter(t, yy, s=1)
    #plt.title('y(t)轨道演化图')
    #plt.xlabel('x轴')
    #plt.ylabel('y轴')
    #plt.savefig('Lorenz的y(t)轨道演化图.png')

    # z(t)轨道演化图
    ##plt.plot(t, z1,  c='purple',linestyle='--' )
    ##plt.plot(t, zz, )
    #plt.scatter(t, z1, s=1, c='purple')
    #plt.scatter(t, zz, s=1)
    #plt.title('z(t)轨道演化图')
    #plt.xlabel('x轴')
    #plt.ylabel('y轴')
    #plt.savefig('Lorenz的z(t)轨道演化图.png')

    plt.show()


if __name__ == '__main__':
    main()
